#include "BeamProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"
#include "VoxelProcess.h"
#include "LocalScene.h"

void BeamProcess::OnInit()
{
	mLocalScene = static_cast<LocalScene*>(mGameScene);
	mVoxelProcess = mLocalScene->GetProcess<VoxelProcess>("VoxelProcess");


	mBeamSize = 8u;
	mBeamWidth = GetBeamSize(GRenderState->GetWidth());
	mBeamHeight = GetBeamSize(GRenderState->GetHeight());


	float tg = glm::tan(mLocalScene->GetCamera().GetFovY() * 0.5f);
	glm::vec3 look = mLocalScene->GetCamera().GetLook();
	glm::vec3 right = mLocalScene->GetCamera().GetRight();
	look = glm::normalize(look);
	right = glm::normalize(right) * tg * mLocalScene->GetCamera().GetAspect();
	glm::vec3 up = glm::normalize(glm::cross(look, right)) * tg;

	CameraUbo cameraUbo;
	cameraUbo.position = glm::vec4(mLocalScene->GetCamera().GetPosition(), 1.f);
	cameraUbo.look = glm::vec4(look, 1.f);
	cameraUbo.right = glm::vec4(right, 1.f);
	cameraUbo.up = glm::vec4(-up, 1.f);



	ENQUEUE_RENDER_COMMAND(BuildModel)([this, cameraUbo](VkCommandBuffer commandBuffer)
		{
			CameraUbo cameraUboData = cameraUbo;
			mCameraUboBuffer = CustomBuffer::Builder()
				.SetData(&cameraUboData)
				.SetSize(sizeof(CameraUbo))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build(this, "cameraUbo");


			RefCountPtr<RHIBuffer> octreeBuffer = GetBuffer("OctreeProcess", "octree");


			mDescriptor = RHIDescriptor::Builder()
				.PushShader<DeferredVertexShader>()
				.PushShader<BeamFragmentShader>()
				.SetDescriptorCount(0, 1)
				.Build();
			mDescriptor->SetValue("octree", octreeBuffer);
			mDescriptor->SetValue("cameraUbo", mCameraUboBuffer);
			mDescriptorMask = mDescriptor->AllocateDescriptor();



			mBeamImage = ColorRenderTarget::Builder()
				.SetWidth(mBeamWidth)
				.SetHeight(mBeamHeight)
				.SetFormat(VK_FORMAT_R32_SFLOAT)
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "beamImage");


			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(mBeamImage)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mBeamImage)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build();

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void BeamProcess::OnTick(float deltaTime)
{

	float tg = glm::tan(mLocalScene->GetCamera().GetFovY() * 0.5f);
	glm::vec3 look = mLocalScene->GetCamera().GetLook();
	glm::vec3 right = mLocalScene->GetCamera().GetRight();
	look = glm::normalize(look);
	right = glm::normalize(right) * tg * mLocalScene->GetCamera().GetAspect();
	glm::vec3 up = glm::normalize(glm::cross(look, right)) * tg;

	CameraUbo cameraUbo;
	cameraUbo.position = glm::vec4(mLocalScene->GetCamera().GetPosition(), 1.f);
	cameraUbo.look = glm::vec4(look, 1.f);
	cameraUbo.right = glm::vec4(right, 1.f);
	cameraUbo.up = glm::vec4(-up, 1.f);


	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, cameraUbo](VkCommandBuffer commandBuffer)
		{
			CameraUbo cameraUboData = cameraUbo;
			mCameraUboBuffer->WriteMemory(&cameraUboData);


			VkViewport viewport{ 0, 0, static_cast<float>(mBeamWidth), static_cast<float>(mBeamHeight), 0.f, 1.f };
			VkRect2D scissor{ 0, 0, mBeamWidth , mBeamHeight };


			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &viewport);
			vkCmdSetScissor(commandBuffer, 0, 1, &scissor);

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mDescriptorMask.GetDescriptorSet(), 0, nullptr);

			PushConstants pushConstants;
			pushConstants.width = GRenderState->GetWidth();
			pushConstants.height = GRenderState->GetHeight();
			pushConstants.beamSize = mBeamSize;
			pushConstants.originSize = exp2f(1.0f - mLocalScene->mOctreeLevel);

			vkCmdPushConstants(commandBuffer, mPipeline->GetPipelineLayout(), VK_SHADER_STAGE_FRAGMENT_BIT, 0, sizeof(PushConstants), &pushConstants);

			vkCmdDraw(commandBuffer, 3, 1, 0, 0);

			vkCmdEndRenderPass(commandBuffer);
		});
}

void BeamProcess::OnResize()
{
	mBeamImage = ColorRenderTarget::Builder()
		.SetWidth(mBeamWidth)
		.SetHeight(mBeamHeight)
		.SetFormat(VK_FORMAT_R32_SFLOAT)
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "beamImage");

	mFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mBeamImage)
		.SetRenderPass(mRenderPass)
		.Build();
}

uint32 BeamProcess::GetBeamSize(uint32 value)
{
	return (value + (mBeamSize - 1)) / mBeamSize + 1;
}
